import subprocess

# 运行nvidia-smi命令获取GPU信息
output = subprocess.check_output(["nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits"])
output_str = output.decode('utf-8')  # 将字节流转换为字符串

# output = output.split
# 将输出按行分割并提取剩余内存值
gpu_memory = [int(x) for x in output_str.strip().split('\n')]

# 打印每个GPU的剩余内存
for i, mem in enumerate(gpu_memory):
    print("GPU {}: 剩余内存 {:.2f} MB".format(i, mem))